#Copyright (C) Nial Peters 2013
#
#This file is part of pydoas.
#
#pydoas is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.
#
#pydoas is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with pydoas.  If not, see <http://www.gnu.org/licenses/>.
import numpy
import types
import os
import datetime
import threading
import seabreeze.backends
seabreeze.backends.use('pyseabreeze')
import seabreeze
import seabreeze.spectrometers

import doas.io
from doas.spec import Spectrum

def list_spectrometers():
    return [d.serial for d in seabreeze.spectrometers.list_devices()]


def get_first_available_spectrometer():
    devices = seabreeze.spectrometers.list_devices()
    if devices:
        return OOISpectrometer(devices[0])
    else:
        return None


class CaptureProgramBase(object):
    def __init__(self, spectrometer):
        
        self._spectrometer = spectrometer
        self._running = False
        self._stay_alive = True
    
    
    @property
    def is_running(self):
        return self._running
    
    
    def initialise(self):
        pass
    
    
    def finalise(self):
        pass
    
    
    def post_process(self, spec):
        pass
    
    
    def _run_capture(self):
        while self._stay_alive:
            st = datetime.datetime.now()
            spec = self._spectrometer.get_spectrum()
            print datetime.datetime.now() - st
            self.post_process(spec)
    
    
    def start(self):
        
        if self._running:
            raise RuntimeError("Capture program is already running")
        self._stay_alive = True
        self._running = True
        self.initialise()
        
        self._worker_thread = threading.Thread(target=self._run_capture)
        
        self._worker_thread.start()
    
    
    def stop(self):
        self._stay_alive = False
        if not self._running:
            return

        self._worker_thread.join()
        self._worker_thread = None
        self.finalise()
        self._running = False
    

    
class CaptureProgramSimple(CaptureProgramBase):
    
    def __init__(self, spectrometer, output_dir, int_time, n_avg):
        self._output_dir = output_dir
        self._int_time = int_time
        self.n_avg = n_avg
        self.__file_writer = doas.io.SpectrumIO()
        
        super(CaptureProgramSimple, self).__init__(spectrometer)
    
    
    def initialise(self):
        self._spectrometer.int_time_millisec = self._int_time
        self._spectrometer.n_avg = self._n_avg
        
        try:
            os.makedirs(self._output_dir)
        except OSError:
            #dir already exists
            pass
    
    
    def post_process(self, spec):
        filename = spec.capture_time.strftime("%Y%m%d_%H%M%S_%%d.txt")
        
        filename = filename%(int(round(spec.capture_time.microsecond /1000.0)))
            
        self.__file_writer.save(spec, os.path.join(self._output_dir, filename), 
                        'SpectraSuite Tab-delimited')
    
    
    def take_darks_and_offsets(self, n=20, dark_int_time=200, offset_n_avg=100):
        
        if self._running:
            raise RuntimeError("Cannot collect darks/offsets while the capture program is running.")
        
        #create the output dirs for the darks
        output_path = os.path.join(self.output_dir,'darks')
        try:
            os.makedirs(output_path)
        except OSError:
            #dir already exists
            pass
        
        self._spectrometer.int_time_millisec = dark_int_time
        self._spectrometer.n_avg = 1
        
        for i in range(n):
            
            spec = self._spectrometer.get_spectrum()
            
            filename = spec.capture_time.strftime("dark_%Y%m%d_%H%M%S_%%f.txt")
        
            filename = filename%(round(spec.capture_time.microsecond * 1000.0))
            
            self.__file_writer.save(spec, os.path.join(output_path, filename), 
                             'SpectraSuite Tab-delimited')
            
        #create the output dirs for the offsets
        output_path = os.path.join(output_dir,'offsets')
        try:
            os.makedirs(output_path)
        except OSError:
            #dir already exists
            pass
        
        self._spectrometer.int_time_millisec = 2
        self._spectrometer.n_avg = offset_n_avg
        
        for i in range(n):
            
            spec = self._spectrometer.get_spectrum()
            
            filename = spec.capture_time.strftime("offset_%Y%m%d_%H%M%S_%%f.txt")
        
            filename = filename%(round(spec.capture_time.microsecond * 1000.0))
            
            self.__file_writer.save(spec, os.path.join(output_path, filename), 
                             'SpectraSuite Tab-delimited')


class CaptureProgramAutoIntTime(CaptureProgramSimple):
    def __init__(self, spectrometer, output_dir, desired_limits=(0.65, 0.85), 
                 wavelength_range=(310, 360), max_int_time=1000, 
                 time_per_spectrum=1000, saturation_level=65535):
        
        super(CaptureProgramAutoIntTime, self).__init__(spectrometer, output_dir,2,1)
 
        self.desired_limits = desired_limits
        self.wavelength_range = wavelength_range
        self.max_int_time = max_int_time
        self.time_per_spectrum = time_per_spectrum
        self.saturation_level = saturation_level
        
        #define the region over which we are going to optimise the integration
        #times
        wavelengths = self._spectrometer.wavelengths
        criteria = numpy.logical_and(wavelengths < self.wavelength_range[1],
                                     wavelengths > self.wavelength_range[0])
     
        self.__optimisation_idxs = numpy.where(criteria)
        
               
        
    def initialise(self):
        self._int_time = self.get_initial_guess()
        self._n_avg = max(1,int(self.time_per_spectrum//self._int_time))
        super(CaptureProgramAutoIntTime, self).initialise()
        
    
    def post_process(self, spectrum):
        #check for saturation
        print "Captured spectrum"
        if numpy.max(spectrum.counts[self.__optimisation_idxs])>=self.saturation_level:
            print "Danger - saturation!"
        
        super(CaptureProgramAutoIntTime, self).post_process(spectrum)
    
    def get_initial_guess(self):
        
        int_times = []
        intensities = []
        
        current_it = self._spectrometer.min_int_time_millisec

        #take spectra at a range of integration times and measure the maximum
        #counts with the optimisation region
        while current_it <= self.max_int_time:
            self._spectrometer.int_time_millisec = int(current_it)
            spectrum = self._spectrometer.get_spectrum()
            
            max_counts = numpy.max(spectrum.counts[self.__optimisation_idxs])
            
            if max_counts >= self.saturation_level:
                break
            
            intensities.append(max_counts)
            int_times.append(current_it)
            current_it *= 2
        
        if len(int_times) < 2:
            return self._spectrometer.min_int_time_millisec
        
        #do a linear fit through the data
        grad,intercept = numpy.polyfit(numpy.array(int_times)[-2:], 
                                       numpy.array(intensities)[-2:], 1)
        level = 0.5 * (self.desired_limits[0] + self.desired_limits[1])
        return int(round(((self.saturation_level * level) - intercept) / grad))


class OOISpectrometer(object):
    def __init__(self, _dev):
        self._ooi_spectro = seabreeze.spectrometers.Spectrometer(_dev)
        
        self.int_time_millisec = int(round((2/1.0e3) * self._ooi_spectro.minimum_integration_time_micros))
        self._ooi_spectro.intensities()
        #print self.int_time_millisec
        #self._ooi_spectro.integration_time_micros(self.__int_time_micros)
        self.__n_avg = 1
        self.__wavelengths = self._ooi_spectro.wavelengths()
    
    @property
    def serial_number(self):
        """
        The serial number of the spectrometer (string).
        """
        return self._ooi_spectro.serial_number
    
    @property
    def int_time_millisec(self):
        """
        The integration time in milliseconds.
        """
        return self.__int_time_micros / 1.0e3
    
    @int_time_millisec.setter
    def int_time_millisec(self, value):
        print "set int time to ", value
        if not type(value) == types.IntType:
            raise TypeError("int_time_millisec must be an integer")
        self.__int_time_micros = value * 1.0e3
        self._ooi_spectro.integration_time_micros(self.__int_time_micros)
        
        #need to collect a spectrum at the new integration time before returning
        #to ensure that the spectrometer has settled
        self._ooi_spectro.intensities()
        
        
    @property
    def wavelengths(self):
        #return self._ooi_spectro.wavelengths()
        return self.__wavelengths
    
    @property
    def min_int_time_millisec(self):
        """
        The minimum integration time that the spectrometer can be set to
        in milliseconds.
        """
        return self._ooi_spectro.minimum_integration_time_micros/1.0e3
    
    @property
    def n_avg(self):
        """
        The number of spectra to average for each returned spectrum.
        """
        return self.__n_avg
    
    @n_avg.setter
    def n_avg(self, value):
        if not type(value) == types.IntType:
            raise TypeError("n_avg must be an integer")
        
        if value < 1:
            raise ValueError("n_avg must be greater than or equal to one")
        
        self.__n_avg = value
    
    
    def get_spectrum(self):
        data = numpy.zeros(self._ooi_spectro.pixels, dtype=numpy.float64)
        capture_time = datetime.datetime.utcnow()
        for i in range(self.__n_avg):
            data += numpy.array(self._ooi_spectro.intensities(), dtype=numpy.float64)
        
        data /= self.__n_avg
        print "int_time = ",self.int_time_millisec
        return Spectrum(None, self.wavelengths, data, capture_time, self.serial_number, 
                        self.int_time_millisec, self.n_avg)
    
    def close(self):
        self._ooi_spectro.close()
        